import sqlite3
import re
from typing import Tuple, Optional
from ask_ollama import *
import csv
from tqdm import tqdm  # 导入进度条库


def get_scenic_info(db_path: str, scenic_id: int) -> Tuple[Optional[str], Optional[str]]:
    """
    从数据库中获取景点全名（去除省份城市前缀）和相关的游记描述（过滤#xxxxx[话题]#格式内容）

    参数:
        db_path: 数据库文件路径
        scenic_id: 景点ID

    返回:
        包含处理后的景点名称和合并的游记描述的元组，如果出错则返回(None, None)
    """
    try:
        # 连接数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # 查询景点全名
        cursor.execute("SELECT full_name FROM scenic_spots WHERE id = ?", (scenic_id,))
        result = cursor.fetchone()

        if not result:
            print(f"未找到ID为{scenic_id}的景点")
            conn.close()
            return (None, None)

        full_name = result[0]

        # 处理景点名称：去除开头的xx省xx市
        pattern = r'^[\u4e00-\u9fa5]+省[\u4e00-\u9fa5]+市'
        cleaned_name = re.sub(pattern, '', full_name).strip()

        # 如果处理后为空，使用原始名称
        if not cleaned_name:
            cleaned_name = full_name

        # 查询该景点的所有游记描述
        cursor.execute("SELECT description FROM travel_notes WHERE scenic_spot_id = ?", (scenic_id,))
        descriptions = cursor.fetchall()
        # 2. 单独查询游记数量（关键修复：新增计数查询）
        cursor.execute("SELECT COUNT(*) FROM travel_notes WHERE scenic_spot_id = ?", (scenic_id,))
        note_count = cursor.fetchone()[0]  # 计数查询必有结果（至少为0），不会返回None
        print(f"查询到ID为{scenic_id}的景点共有{note_count}篇游记")

        # 定义正则表达式匹配#xxxxx[话题]#格式的内容
        topic_pattern = r'#.*?\[话题\]#'

        # 处理每条描述，过滤指定格式内容后再合并
        processed_descriptions = []
        for desc in descriptions:
            if desc[0]:
                # 移除匹配的内容
                cleaned_desc = re.sub(topic_pattern, '', desc[0])
                # 去除可能产生的多余空白
                cleaned_desc = cleaned_desc.strip()
                if cleaned_desc:  # 只添加非空内容
                    processed_descriptions.append(cleaned_desc)

        # 合并所有处理后的描述
        merged_description = "\n\n".join(processed_descriptions)

        # 关闭连接
        conn.close()

        return (cleaned_name, merged_description)

    except sqlite3.Error as e:
        print(f"数据库操作错误: {e}")
        return (None, None)
    except Exception as e:
        print(f"发生错误: {e}")
        return (None, None)


if __name__ == "__main__":
    db_path = "./jiangsu2.db"
    # output_file = "scenic_spots_introduction.csv"
    # total_scenic = 54  # 总景点数量，用于进度条
    #
    # # 1. 准备CSV文件并写入表头：新增“景点ID”列作为第一列（顺序可调整）
    # with open(output_file, mode='w', encoding='utf-8', newline='') as file:
    #     writer = csv.writer(file)
    #     # 修改：表头添加“景点ID”，顺序可根据需求调整（如放在最后一列）
    #     writer.writerow(["景点ID", "景点名称", "简介"])
    #
    #     # 使用tqdm创建进度条，遍历所有景点ID
    #     for scenic_id in tqdm(range(1, total_scenic + 1), desc="处理进度", unit="个景点"):
    #         name, description = get_scenic_info(db_path, scenic_id)
    #
    #         if name:
    #             question = f"{description}请你根据以上信息生成“{name}”的景点简介，要求整理成一段即可，不需要额外整理景点信息，简介即可"
    #             introduction = get_answer(question)
    #
    #             # 2. 写入每行数据：补充scenic_id（与表头顺序对应）
    #             # 修改：每行数据添加“景点ID”，位置与表头“景点ID”列对应
    #             writer.writerow([scenic_id, name, introduction])
    #
    # print(f"所有景点信息已成功保存到 {output_file}")

    name, description = get_scenic_info(db_path, 18)
    question = f"{description}请你根据以上信息生成“{name}”的景点简介，要求整理成一段即可，不需要额外整理景点信息，简介即可"
    print(get_answer(question))
